/*
 * CFileMediaSink.cpp
 *
 *  Created on: 2016年1月22日
 *      Author: terry
 */

#include "CFileMediaSink.h"
#include "TFileUtil.h"
#include "CLog.h"
#include "AacHelper.h"
#include "Path.h"
#include "TimeHelper.h"
#include "FfmpegUtil.h"
#include "H264PropParser.h"


namespace av
{

static const int MAX_PKT_SIZE = 100 * 1024;

CFileMediaSink::CFileMediaSink():
		m_fmt(),
		m_fileName(),
		m_maxDuration(),
		m_videoIndex(-1),
		m_audioIndex(-1),
		m_videoStartTime(),
		m_audioStartTime(),
		m_fmtCtx(),
        m_changeFile(),
        m_doAudioTranscode(),
		m_videoPktCount()
{
    m_audioBuffer.ensure(MAX_PKT_SIZE);
}

CFileMediaSink::~CFileMediaSink()
{
	close();
}

bool CFileMediaSink::open(const char* filename)
{
    if (isOpen())
    {
        close();
    }

    m_fileName = filename;

    start();

    return true;
}

void CFileMediaSink::close()
{
    if (isRunning())
    {
        stop();
    }

    closeFile();
}

bool CFileMediaSink::isOpen()
{
    return isRunning();
}

bool CFileMediaSink::openFile(const char* filename, const MediaFormat& fmt)
{
    if (strlen(filename) <= 0)
    {
        return false;
    }

    /// 还没有设置输出分辨率
    if (!m_fmt.isValid())
    {
        return false;
    }

    if (!openFile())
    {
    	return false;
    }

    return true;
}

const char* CFileMediaSink::getFile()
{
	return m_fileName.c_str();
}

int CFileMediaSink::run()
{
    while (!m_canExit)
    {
        AVPacketPtr pkt;
        m_pktQueue.pop(pkt, -1);
        if (pkt)
        {
            handlePacket(pkt);
        }
    }
    return 0;
}

void CFileMediaSink::doStop()
{
    m_pktQueue.cancelWait();
}


int CFileMediaSink::setFile(const char* filename)
{
	CLog::info("CFileMediaSink::setFile: %s\n", filename);

	if (filename == NULL)
	{
		m_fileName.clear();
	}
	else
	{
		m_fileName = filename;
	}

    {
	    comn::AutoCritSec lock(m_cs);
	    if (isFileOpen())
	    {
		    m_changeFile = true;
	    }
    }
    
    m_pktQueue.cancelWait();

	return 0;
}

void CFileMediaSink::onMediaFormat(const MediaFormat& fmt)
{
	m_fmt = fmt;

	if (m_fmt.m_videoProp.size() > 0)
	{
		//comn::FileUtil::write(m_fmt.m_videoProp.data(), m_fmt.m_videoProp.size(), "H:\\dump\\prop.bin", false);
	}

	m_fmt.m_videoProp.clear(); // debug
}

void CFileMediaSink::onMediaPacket(AVPacketPtr& pkt)
{
	if (!pkt)
	{
		return;
	}

    m_pktQueue.push(pkt);
}

void CFileMediaSink::onMediaEvent(int event)
{
    // pass
}

bool CFileMediaSink::openFile()
{
	int rc = avformat_alloc_output_context2(&m_fmtCtx, NULL, NULL, m_fileName.c_str());
	if (!m_fmtCtx)
	{
		return false;
	}

	if (m_fmt.hasVideo())
	{
		AVCodecID codecID = (AVCodecID)m_fmt.m_codec;
		AVCodec* codec = avcodec_find_decoder(codecID);
		if (codec != NULL)
		{
			AVStream *out_stream = avformat_new_stream(m_fmtCtx, NULL);
			AVCodecContext * codecCtx = out_stream->codec;

			avcodec_get_context_defaults3(codecCtx, codec);
            
            if (m_fmt.m_width <= 0)
            {
                codecCtx->width = 1920;
                codecCtx->height = 1080;
            }
            else
            {
			    codecCtx->width = m_fmt.m_width;
			    codecCtx->height = m_fmt.m_height;
			    codecCtx->profile = m_fmt.m_profile;
            }

			codecCtx->time_base = av_make_q(1, 90000);

			out_stream->time_base = codecCtx->time_base;

			// extra_data
			if (m_fmt.m_videoProp.size() > 0)
			{
				codecCtx->extradata = (uint8_t*)av_memdup(m_fmt.m_videoProp.c_str(), m_fmt.m_videoProp.size());
				codecCtx->extradata_size = m_fmt.m_videoProp.size();
			}

			if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
			{
				out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
			}

			m_videoIndex = 0;
		}
	}

	if (m_fmt.hasAudio())
	{
		AVCodecID codecID = (AVCodecID)m_fmt.m_audioCodec;
		if (isG711(m_fmt) && isMp4(m_fmtCtx))
		{
			codecID = AV_CODEC_ID_NONE;

            //m_doAudioTranscode = true;
			//m_audioTranscoder.open(m_fmt, codecID, m_fmt.m_channels, m_fmt.m_sampleRate);
		}

		AVCodec* codec = avcodec_find_encoder(codecID);
		if (codec != NULL)
		{
			AVStream *out_stream = avformat_new_stream(m_fmtCtx, codec);
			AVCodecContext * codecCtx = out_stream->codec;

			avcodec_get_context_defaults3(codecCtx, codec);

            if (m_audioTranscoder.isOpen())
            {
			    avcodec_copy_context(codecCtx, m_audioTranscoder.getCodecContext());

                out_stream->time_base = codecCtx->time_base;
                if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
                {
                    out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
                }

            }
            else
            {
                codecCtx->channels = m_fmt.m_channels;
                codecCtx->channel_layout = av_get_default_channel_layout(m_fmt.m_channels);
                codecCtx->sample_rate = m_fmt.m_sampleRate;
                codecCtx->time_base.num = 1;
                codecCtx->time_base.den = m_fmt.m_sampleRate;
                codecCtx->frame_size = m_fmt.m_frameSize;

                out_stream->time_base = codecCtx->time_base;

                if (m_fmt.m_audioConfig.size() > 0)
                {
                    codecCtx->extradata = (uint8_t*)av_memdup(m_fmt.m_audioConfig.c_str(), m_fmt.m_audioConfig.size());
                    codecCtx->extradata_size = m_fmt.m_audioConfig.size();
                }

                if (m_fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
                {
                    out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
                }
            }

			m_audioIndex = m_videoIndex + 1;
		}
	}


	rc = avio_open(&m_fmtCtx->pb, m_fileName.c_str(), AVIO_FLAG_WRITE);
	if (m_fmtCtx->pb == NULL)
	{
		avformat_free_context(m_fmtCtx);
		m_fmtCtx = NULL;
        m_videoIndex = -1;
        m_audioIndex = -1;
		return false;
	}

    rc = avformat_write_header(m_fmtCtx, NULL);

	return true;
}

void CFileMediaSink::closeFile()
{
    if (m_fmtCtx)
    {
    	av_write_trailer(m_fmtCtx);

    	avio_close(m_fmtCtx->pb);

        avformat_free_context(m_fmtCtx);
        m_fmtCtx = NULL;
    }

    m_audioTranscoder.close();

    m_changeFile = false;

    m_videoIndex = -1;
    m_audioIndex = -1;
    m_videoStartTime = 0;
    m_audioStartTime = 0;
}

bool CFileMediaSink::isFileOpen()
{
	return (m_fmtCtx != NULL);
}

bool CFileMediaSink::isG711(const MediaFormat& fmt)
{
	return (fmt.m_audioCodec == AV_CODEC_ID_PCM_ALAW) || (fmt.m_audioCodec == AV_CODEC_ID_PCM_MULAW);
}

bool CFileMediaSink::isMp4(AVFormatContext* fmtCtx)
{
	int order = strcmp(fmtCtx->oformat->name, "mp4");
	return (order == 0);
}

void CFileMediaSink::writeVideo(AVPacketPtr& inPkt)
{
    if (m_videoIndex < 0)
    {
        return;
    }

	if ((m_videoPktCount == 0) && ((inPkt->flags & AV_PKT_FLAG_KEY) == 0))
	{
		/// 第一包必须是关键帧
		return;
	}
	m_videoPktCount++;

	if ((m_videoStartTime == 0) || (m_videoStartTime == AV_NOPTS_VALUE))
	{
		m_videoStartTime = inPkt->dts;
	}
    else if (inPkt->dts < m_videoStartTime)
    {
        m_videoStartTime = inPkt->dts;
    }

	int64_t dts = inPkt->dts - m_videoStartTime;

	AVPacket pkt;
	av_init_packet(&pkt);

	av_packet_ref(&pkt, inPkt.get());

	pkt.stream_index = m_videoIndex;
	pkt.pts = dts;
	pkt.dts = dts;

    AVRational src = av_make_q(1, AV_TIME_BASE);
    AVRational dst = av_make_q(1, 90000);
    av_packet_rescale_ts(&pkt, src, dst);

    //CLog::debug("video packet. size:%d, duration:%d, dts:%I64d, pts:%I64d\n", pkt.size, pkt.duration/90, pkt.dts/90, pkt.pts/90);

	av_interleaved_write_frame(m_fmtCtx, &pkt);

    av_packet_unref(&pkt);
}

void CFileMediaSink::writeAudio(AVPacketPtr& inPkt)
{
	//av_usleep(1000);

	//if (false)
	//{
	//	if (!m_audioInPoint.isSet())
	//	{
	//		m_audioInPoint.reset(util::TimeHelper::getTime(), inPkt->dts);
	//		CLog::debug("input doWriteAudio packet. clock:%I64d, dts:%I64d\n", m_audioInPoint.m_clock, m_audioInPoint.m_pts);
	//	}

	//	TimePoint curPoint(util::TimeHelper::getTime(), inPkt->dts);
	//	int clockDuration = (int)((curPoint.m_clock - m_audioInPoint.m_clock) / 1000);
	//	int ptsDuration = (int)((curPoint.m_pts - m_audioInPoint.m_pts) / 1000);
	//	int dist = clockDuration - ptsDuration;

	//	//if (abs(dist) > 100)
	//	{
	//		//CLog::debug("input audio packet. clock duration:%d, pts duration:%d. dist:%d\n", clockDuration, ptsDuration, dist);
	//	}
	//}

    if (m_doAudioTranscode)
    {
        AVPacket outPkt;
        av_init_packet(&outPkt);
        outPkt.data = m_audioBuffer.data();
        outPkt.size = m_audioBuffer.capacity();

        if (m_audioTranscoder.transcode(inPkt.get(), &outPkt))
        {
            outPkt.stream_index = m_audioIndex;
            doWriteAudio(&outPkt);
            
            while (true)
            {
                av_init_packet(&outPkt);
                outPkt.data = m_audioBuffer.data();
                outPkt.size = m_audioBuffer.capacity();

                if (!m_audioTranscoder.transcode(NULL, &outPkt))
                {
                    break;
                }

                outPkt.stream_index = m_audioIndex;
                doWriteAudio(&outPkt);
            }
        }
    }
    else
    {
        if (m_audioStartTime == 0)
        {
            m_audioStartTime = inPkt->pts;
        }

        int64_t pts = inPkt->pts - m_audioStartTime;
        

        AVPacket pkt;
        av_init_packet(&pkt);

        av_packet_ref(&pkt, inPkt.get());
        pkt.stream_index = m_audioIndex;
        pkt.pts = pts;
        pkt.dts = pts;

        AVRational src = av_make_q(1, AV_TIME_BASE);
        AVRational dst = av_make_q(1, m_fmt.m_sampleRate);
        av_packet_rescale_ts(&pkt, src, dst);

        doWriteAudio(&pkt);

        av_packet_unref(&pkt);
    }
}

void CFileMediaSink::doWriteAudio(AVPacket* pkt)
{
	//if (!m_audioPoint.isSet())
	//{
	//	m_audioPoint.reset(util::TimeHelper::getTime(), pkt->dts);
	//	CLog::debug("doWriteAudio packet. clock:%I64d, dts:%I64d\n", m_audioPoint.m_clock, m_audioPoint.m_pts);
	//}

	//TimePoint curPoint(util::TimeHelper::getTime(), pkt->dts);
	//int clockDuration = (int)((curPoint.m_clock - m_audioPoint.m_clock) / 1000);
	//int ptsDuration = (int)((curPoint.m_pts - m_audioPoint.m_pts) * 1000 / m_fmt.m_sampleRate);
	//int dist = clockDuration - ptsDuration;

	//if (abs(dist) > 1000)
	//{
        //int64_t dts = av_rescale(pkt->dts, 1000, m_fmt.m_sampleRate);
        //int64_t pts = av_rescale(pkt->pts, 1000, m_fmt.m_sampleRate);
        //int64_t dts = pkt->dts;
        //int64_t pts = pkt->pts;
        //CLog::debug("audio packet. size:%d, duration:%d, dts:%I64d, pts:%I64d\n", pkt->size, pkt->duration, dts, pts);
	//}

    av_interleaved_write_frame(m_fmtCtx, pkt);
}

void CFileMediaSink::handlePacket(AVPacketPtr& pkt)
{
    //assert(pkt->size > 0);
	if (m_fmt.m_videoProp.empty())
	{
		if (pkt->stream_index != 0)
		{
			return;
		}

		if (!parseVideoProp(pkt))
		{
			return;
		}
	}

    if (m_changeFile)
    {
		closeFile();
    }

    if (!m_fileName.empty() && m_fmt.isValid())
    {
        if (!isFileOpen())
        {
            openFile(m_fileName.c_str(), m_fmt);
        }
    }

    if (pkt->stream_index == 0)
    {
        writeVideo(pkt);
    }
    else if (pkt->stream_index == m_audioIndex)
    {
        writeAudio(pkt);
    }
}

bool CFileMediaSink::parseVideoProp(AVPacketPtr& pkt)
{
	av::NaluPacket nalu;
	if (!H264PropParser::parseNalu(pkt->data, pkt->size, nalu))
	{
		return false;
	}

	if (nalu.type != av::NaluPacket::NALU_SPS)
	{
		return false;
	}

	std::string sps;
	std::string pps;
	if (!H264PropParser::splitPropSet(pkt->data, pkt->size, sps, pps))
	{
		return false;
	}

	std::string sprop;
	sprop.append((char*)H264PropParser::getStartCode(), H264PropParser::START_CODE_LENGTH);
	sprop.append(sps);
	sprop.append((char*)H264PropParser::getStartCode(), H264PropParser::START_CODE_LENGTH);
	sprop.append(pps);
	//H264PropParser::combineToSprop(sps, pps, sprop);

	m_fmt.m_videoProp.assign(sprop.c_str(), sprop.size());

	//comn::FileUtil::write(m_fmt.m_videoProp.data(), m_fmt.m_videoProp.size(), "H:\\dump\\prop_parse.bin", false);

	return true;
}



} /* namespace av */
